In [4]:
#Need to uncomment and run this when the container is first instantiated
#!apt-get install python3-opencv
In [5]:
import cv2 as cv
In [6]:
from PIL import Image
import sys
import os
import urllib
import tensorflow.contrib.tensorrt as trt
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import tensorflow as tf
import numpy as np
import time
from tf_trt_models.detection import download_detection_model, build_detection_graph
%matplotlib inline
Download an SSD model for face detection¶
In [7]:
# https://github.com/yeephycho/tensorflow-face-detection
FROZEN_GRAPH_NAME = 'data/frozen_inference_graph_face.pb'
!wget https://github.com/yeephycho/tensorflow-face-detection/blob/master/model/frozen_inference_graph_face.pb?raw=true -O {FROZEN_GRAPH_NAME}
In [73]:
cv.imwrite('data/my_sample_webcam_pic.png',img)
Out[73]:
Download an image with some faces¶
In [8]:
# # go warriors!
# IMAGE_PATH = 'data/warriors.jpg'
# !wget 'https://cdn.vox-cdn.com/thumbor/rC0mlBATZdoDW1tEa44P6431sGc=/0x0:3683x2455/1200x800/filters:focal(1623x234:2211x822)/cdn.vox-cdn.com/uploads/chorus_image/image/63273148/usa_today_12005182.0.jpg' -O {IMAGE_PATH}
Load the frozen graph¶
In [9]:
output_dir=''
frozen_graph = tf.GraphDef()
with open(os.path.join(output_dir, FROZEN_GRAPH_NAME), 'rb') as f:
frozen_graph.ParseFromString(f.read())
A few magical constants¶
In [10]:
# https://github.com/NVIDIA-AI-IOT/tf_trt_models/blob/master/tf_trt_models/detection.py
INPUT_NAME='image_tensor'
BOXES_NAME='detection_boxes'
CLASSES_NAME='detection_classes'
SCORES_NAME='detection_scores'
MASKS_NAME='detection_masks'
NUM_DETECTIONS_NAME='num_detections'
input_names = [INPUT_NAME]
output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]
Optimize the frozen graph using TensorRT¶
In [11]:
trt_graph = trt.create_inference_graph(
input_graph_def=frozen_graph,
outputs=output_names,
max_batch_size=1,
max_workspace_size_bytes=1 << 25,
precision_mode='FP16',
minimum_segment_size=50
)
Create session and load graph¶
In [12]:
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
tf_sess = tf.Session(config=tf_config)
# use this if you want to try on the optimized TensorRT graph
# Note that this will take a while
# tf.import_graph_def(trt_graph, name='')
# use this if you want to try directly on the frozen TF graph
# this is much faster
tf.import_graph_def(frozen_graph, name='')
tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')
tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')
Capture Images & Run Inference¶
In [23]:
# suppress boxes that are below the threshold..
DETECTION_THRESHOLD = 0.5
In [50]:
def get_images(cap):
_ret, img = cap.read()
if not _ret:
print('No image')
return
img_color = cv.cvtColor(img, cv.COLOR_BGR2RGB)
#print("The image as obtained from the webcam")
#plt.imshow(img_color)
#plt.show()
image = Image.fromarray(img_color)
image_resized = np.array(image.resize((300, 300)))
#image_resized = np.resize(image, (300,300))
#print(image_resized.shape)
#print("The image that has been resized to run inference")
#plt.imshow(image_resized)
#plt.show()
image = np.array(image)
run_inference(image_resized, image)
def run_inference(image_resized, image):
scores, boxes, classes, num_detections = tf_sess.run(
[tf_scores, tf_boxes, tf_classes, tf_num_detections],
feed_dict={ tf_input: image_resized[None, ...] })
boxes = boxes[0] # index by 0 to remove batch dimension
scores = scores[0]
classes = classes[0]
num_detections = num_detections[0]
display_detected_images(image, boxes, scores, num_detections, classes)
def display_detected_images(image, boxes, scores, num_detections, classes):
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.imshow(image)
# plot boxes exceeding score threshold
for i in range(int(num_detections)):
if scores[i] < DETECTION_THRESHOLD:
continue
# scale box to image coordinates
box = boxes[i] * np.array([image.shape[0], image.shape[1], image.shape[0], image.shape[1]])
# display rectangle
patch = patches.Rectangle((box[1], box[0]), box[3] - box[1], box[2] - box[0], color='g', alpha=0.3)
ax.add_patch(patch)
# display class index and score
plt.text(x=box[1] + 10, y=box[2] - 10, s='%d (%0.2f) ' % (classes[i], scores[i]), color='w')
print("Image with face being detected with the box")
plt.show()
In [57]:
cap = cv.VideoCapture(1)
In [58]:
for i in range(30):
get_images(cap)
In [56]:
cap.release()
Benchmark¶
In [81]:
num_samples = 50
t0 = time.time()
for i in range(num_samples):
scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
tf_input: image_resized[None, ...]
})
t1 = time.time()
print('Average runtime: %f seconds' % (float(t1 - t0) / num_samples))
Close session to release resources¶
In [16]:
tf_sess.close()
In [ ]: